import re
from functools import cached_property
from typing import ClassVar
import itertools
from typing_extensions import Self
from pydantic import BaseModel, model_validator
from pysam import FastaFile  # pylint: disable=E0611
from .load_data import BEDFrame
from .transcript import Interval, Transcript
from .sequence import DonorSequence, AcceptorSequence, Sequence
from .snv import SNV


class Splicing(BaseModel):
    DonorThreshold: ClassVar[float] = 3
    AcceptorThreshold: ClassVar[float] = 3
    PercentThreshold: ClassVar[float] = 0.7
    DonorExon: ClassVar[int] = 3
    DonorIntron: ClassVar[int] = 6
    AcceptorExon: ClassVar[int] = 3
    AcceptorIntron: ClassVar[int] = 20
    GenomeFasta: ClassVar[FastaFile]
    Domain: ClassVar[BEDFrame]
    Hotspot: ClassVar[BEDFrame]
    CuratedRegion: ClassVar[BEDFrame]

    snv: SNV
    transcript: Transcript
    start: int = 0
    end: int = 0
    ref: DonorSequence | AcceptorSequence = DonorSequence("")
    alt: DonorSequence | AcceptorSequence = DonorSequence("")
    type: str = "NA"
    index: str = "NA"

    class Config:
        arbitrary_types_allowed = True

    @cached_property
    def ref_maxent_score(self) -> float:
        return self.ref.score(-1.0)

    @cached_property
    def alt_maxent_score(self) -> float:
        return self.alt.score(-1.0)

    @cached_property
    def maxent_foldchange(self) -> float:
        return round(self.ref_maxent_score / self.alt_maxent_score, 2)

    def _parse_forward_donor(self, intron_start: int) -> tuple[str, int, int, DonorSequence, DonorSequence]:
        exon_alt = intron_alt = ""
        start = intron_start - self.DonorExon
        end = intron_start + self.DonorIntron
        ref = self.GenomeFasta.fetch(self.snv.chrom, start, end)
        if intron_start - self.DonorExon < self.snv.pos - 1:
            exon_alt = self.GenomeFasta.fetch(self.snv.chrom, start, self.snv.pos - 1)
        if self.snv.pos + len(self.snv.ref) - 1 < end + len(self.snv.ref) - len(self.snv.alt):
            intron_alt = self.GenomeFasta.fetch(self.snv.chrom, self.snv.pos + len(self.snv.ref) - 1, end + len(self.snv.ref) - len(self.snv.alt))
        alt = exon_alt + self.snv.alt + intron_alt
        return "donor", start, end, DonorSequence(ref), DonorSequence(alt)

    def _parse_forward_acceptor(self, intron_end: int) -> tuple[str, int, int, AcceptorSequence, AcceptorSequence]:
        exon_alt = intron_alt = ""
        start = intron_end - self.AcceptorIntron
        end = intron_end + self.AcceptorExon
        ref = self.GenomeFasta.fetch(self.snv.chrom, start, end)
        if self.snv.pos + len(self.snv.ref) - 1 < end:
            exon_alt = self.GenomeFasta.fetch(self.snv.chrom, self.snv.pos + len(self.snv.ref) - 1, end)
        if start - len(self.snv.ref) + len(self.snv.alt) < self.snv.pos - 1:
            intron_alt = self.GenomeFasta.fetch(self.snv.chrom, start - len(self.snv.ref) + len(self.snv.alt), self.snv.pos - 1)
        altseq = intron_alt + self.snv.alt + exon_alt
        return "acceptor", start, end, AcceptorSequence(ref), AcceptorSequence(altseq)

    def _parse_backward_acceptor(self, intron_start: int) -> tuple[str, int, int, AcceptorSequence, AcceptorSequence]:
        exon_alt = intron_alt = ""
        start = intron_start - self.AcceptorExon
        end = intron_start + self.AcceptorIntron
        # ref = Sequence(self.genome.fetch(self.snv.chrom, start, end)).reverse_complement()
        ref = self.GenomeFasta.fetch(self.snv.chrom, start, end)
        if start < self.snv.pos - 1:
            exon_alt = self.GenomeFasta.fetch(self.snv.chrom, start, self.snv.pos - 1)
        if self.snv.pos + len(self.snv.ref) - 1 < end + len(self.snv.ref) - len(self.snv.alt):
            intron_alt = self.GenomeFasta.fetch(self.snv.chrom, self.snv.pos + len(self.snv.ref) - 1, end + len(self.snv.ref) - len(self.snv.alt))
        alt = exon_alt + self.snv.alt + intron_alt
        return ("acceptor", start, end, AcceptorSequence(ref).reverse_complement(), AcceptorSequence(alt).reverse_complement())

    def _parse_backward_donor(self, intron_end: int) -> tuple[str, int, int, DonorSequence, DonorSequence]:
        exon_alt = intron_alt = ""
        start = intron_end - self.DonorIntron
        end = intron_end + self.DonorExon
        ref = self.GenomeFasta.fetch(self.snv.chrom, start, end)
        if self.snv.pos + len(self.snv.ref) - 1 < end:
            exon_alt = self.GenomeFasta.fetch(self.snv.chrom, self.snv.pos + len(self.snv.ref) - 1, end)
        if start - len(self.snv.ref) + len(self.snv.alt) < self.snv.pos - 1:
            intron_alt = self.GenomeFasta.fetch(
                self.snv.chrom,
                start - len(self.snv.ref) + len(self.snv.alt),
                self.snv.pos - 1,
            )
        alt = intron_alt + self.snv.alt + exon_alt
        return ("donor", start, end, DonorSequence(ref).reverse_complement(), DonorSequence(alt).reverse_complement())

    @model_validator(mode="after")
    def parse(self) -> Self:
        # pylint: disable=W0640
        """
        Get the refseq and altseq around splice sites
        """
        intron_count = self.transcript.exon_count - 1
        start, end, ref, alt, index, typo = (0, 0, DonorSequence(""), DonorSequence(""), "NA", "NA")
        for i, intron in enumerate(self.transcript.intron_intervals):
            for to_start, to_end in map(
                lambda x: (x - intron.start, x - intron.end),
                range(self.snv.pos, self.snv.pos + len(self.snv.ref)),
            ):
                if self.transcript.strand == "+":
                    if -self.DonorExon < to_start <= self.DonorIntron:
                        typo, start, end, ref, alt = self._parse_forward_donor(intron.start)
                        index = f"IVS{i+1}+{to_start}" if to_start > 0 else f"EX{i+1}-{1-to_start}"
                    if -self.AcceptorIntron < to_end <= self.AcceptorExon:
                        typo, start, end, ref, alt = self._parse_forward_acceptor(intron.end)
                        index = f"EX{i+2}+{to_end}" if to_end > 0 else f"IVS{i+1}-{1-to_end}"
                else:
                    if -self.AcceptorExon < to_start <= self.AcceptorIntron:
                        typo, start, end, ref, alt = self._parse_backward_acceptor(intron.start)
                        index = f"IVS{intron_count-i}-{to_start}" if to_start > 0 else f"EX{intron_count - i + 1}+{1-to_start}"
                    if -self.DonorIntron < to_end <= self.DonorExon:
                        typo, start, end, ref, alt = self._parse_backward_donor(intron.end)
                        index = f"EX{intron_count - i}-{to_end}" if to_end > 0 else f"IVS{intron_count - i}-{1-to_end}"
        self.index, self.type, self.start, self.end, self.ref, self.alt = (index, typo, start, end, ref, alt)
        return self

    @cached_property
    def cryptic_splice_site(
        self,
    ) -> tuple[int, DonorSequence | AcceptorSequence, float]:
        search_flank = 50
        search_region = list(
            itertools.chain.from_iterable(
                zip(
                    range(self.start - 1, self.start - 1 - search_flank, -1),
                    range(self.start + 1, self.start + 1 + search_flank, 1),
                )
            )
        )
        for pos in search_region:
            if self.type == "donor":
                context = DonorSequence(self.GenomeFasta.fetch(self.snv.chrom, pos, pos + 9))
                index = self.snv.pos - pos - 1
                if 0 < index < 9:
                    context = context[:index] + self.snv.alt + context[index + len(self.snv.alt) : 10 - len(self.snv.alt)]
                if self.transcript.strand == "-":
                    context = context.reverse_complement()
                score = context.score(default=0)
                if context[3:5] in ["GT", self.ref[3:5]] and score > 1 and (score >= self.DonorThreshold or score / self.ref_maxent_score >= self.PercentThreshold):
                    return pos, context, score
            elif self.type == "acceptor":
                context = AcceptorSequence(self.GenomeFasta.fetch(self.snv.chrom, pos, pos + 23))
                index = self.snv.pos - pos - 1
                if 0 < index < 23:
                    context = context[:index] + self.snv.alt + context[index + len(self.snv.alt) : 24 - len(self.snv.alt)]
                if self.transcript.strand == "-":
                    context = context.reverse_complement()
                score = context.score(default=0)
                if context[18:20] in ["AG", self.ref[18:20]] and score > 1 and (score >= self.AcceptorThreshold or score / self.ref_maxent_score >= self.PercentThreshold):
                    return pos, context, score
        return 0, DonorSequence(""), 0

    @cached_property
    def skipped_exon_id(self) -> int:
        m = re.match(r"IVS(\d+)([+|-])(\d+)", self.index)
        if m:
            return int(m.group(1)) if m.group(2) == "+" else (int(m.group(1)) + 1)
        m = re.match(r"EX(\d+)([+|-])(\d+)", self.index)
        return int(m.group(1)) if m else 0

    @property
    def cryptic_coding_exons(self) -> list[Interval | None]:
        intervals = self.transcript.coding_exon_intervals.copy()
        interval = intervals[self.skipped_exon_id - 1]
        if interval is not None:
            if (self.transcript.strand == "+" and self.type == "donor") or (self.transcript.strand == "-" and self.type == "acceptor"):
                interval.end = self.cryptic_splice_site[0] + 3
            elif self.type == "acceptor":
                interval.start = self.cryptic_splice_site[0] + 20
            else:
                interval.start = self.cryptic_splice_site[0] + 6
            intervals[self.skipped_exon_id - 1] = interval
        return intervals

    @cached_property
    def has_cryptic_splice_site(self) -> bool:
        if self.type != "NA" and self.cryptic_splice_site[0] > 0:
            cryptic_exon = self.cryptic_coding_exons[self.skipped_exon_id - 1]
            if cryptic_exon and cryptic_exon.start < cryptic_exon.end:
                return True
        return False

    @cached_property
    def is_exon_skipping(self) -> bool:
        return not self.has_cryptic_splice_site and self.maxent_foldchange < self.PercentThreshold and self.alt_maxent_score < 3

    @cached_property
    def skipped_exon_size(self) -> int:
        m = re.match(r"IVS(\d+)([+|-])(\d+)", self.index)
        if m:
            intron_id = (int(m.group(1)) - 1) if m.group(2) == "+" else int(m.group(1))
            return self.transcript.coding_exon_sizes[intron_id]
        m = re.match(r"EX(\d+)([+|-])(\d+)", self.index)
        return self.transcript.coding_exon_sizes[int(m.group(1)) - 1] if m else 0

    @cached_property
    def preserves_reading_frame(self) -> bool:
        if self.is_exon_skipping:
            return self.skipped_exon_size % 3 == 0
        if self.has_cryptic_splice_site:
            return (self.cryptic_splice_site[0] - self.start) % 3 == 0
        return True

    @cached_property
    def trans_seq_info(self) -> tuple[Sequence, int, bool]:
        trans_seq = Sequence("")
        cds_sizes = []
        for exon in self.cryptic_coding_exons:
            cds_sizes.append(exon.end - exon.start if exon else 0)
            seq = Sequence(self.GenomeFasta.fetch(exon.chrom, exon.start, exon.end) if exon else "")
            trans_seq += seq if self.transcript.strand == "+" else seq.reverse_complement()
        stop_codon = 0
        for pos in range(0, len(trans_seq), 3):
            if trans_seq[pos : pos + 3].upper() in ["TAA", "TAG", "TGA"]:
                stop_codon = pos
                break
        is_nmd_target = False
        if len(cds_sizes) == 1 or len([i for i in cds_sizes if i > 0]) == 1:
            is_nmd_target = True
        else:
            nmd_cutoff = sum(cds_sizes[:-1]) - min(50, cds_sizes[-2])
            if stop_codon <= nmd_cutoff:
                is_nmd_target = True
        return trans_seq, stop_codon, is_nmd_target

    @cached_property
    def is_undergo_nmd(self) -> bool:
        if self.preserves_reading_frame:
            return False
        if self.has_cryptic_splice_site:
            return self.trans_seq_info[2]
        return self.skipped_exon_id > len(self.transcript.exon_starts)

    @cached_property
    def is_critical_to_protein_func(self) -> bool:
        """
        Truncated/altered region is critical to protein function.
        """
        interval = self.transcript.exon_intervals[self.skipped_exon_id - 1]
        if self.has_cryptic_splice_site and self.preserves_reading_frame:
            start, end = (interval.start, self.cryptic_splice_site[0]) if self.type == "acceptor" else (self.cryptic_splice_site[0], interval.end)
            if start >= end:
                return False
        elif self.has_cryptic_splice_site and not self.preserves_reading_frame:
            start, end = (self.cryptic_splice_site[0], self.transcript.cds_end) if self.transcript.strand == "+" else (self.transcript.cds_start, self.cryptic_splice_site[0])
        elif self.is_exon_skipping:
            start, end = interval.start, interval.end
        else:
            return False
        domain_row = self.Domain.fetch_row(self.snv.chrom, start, end)
        hotspot_row = self.Hotspot.fetch_row(self.snv.chrom, start, end)
        curated_region_row = self.CuratedRegion.fetch_row(self.snv.chrom, start, end)
        if not (curated_region_row or hotspot_row):
            if domain_row:
                missense_plp, missense_blb = domain_row["name"].split("|")[5:7]
                return (int(missense_blb) == 0 and int(missense_plp) >= 5) or (int(missense_blb) > 0 and int(missense_plp) / int(missense_blb) >= 10)
            return False
        return True

    @cached_property
    def variant_removes_10_percent_of_protein(self) -> bool:
        size = abs(self.start - self.cryptic_splice_site[0]) if self.has_cryptic_splice_site else self.skipped_exon_size
        return size / self.transcript.cds_size > 0.1
